# visualize open loop trajectory
import numpy as np
import open3d as o3d
from scipy.spatial.transform import Rotation as R


def create_grid(grid_size=2.0, grid_interval=0.1):
    """
    创建一个XY平面的网格 (LineSet)
    grid_size: 网格边长
    grid_interval: 格子间距
    """
    half = grid_size / 2.0
    lines = []
    points = []

    # 生成网格点
    for i in np.arange(-half, half + grid_interval, grid_interval):
        # 垂直线 (x = const)
        points.append([i, -half, 0])
        points.append([i, half, 0])
        lines.append([len(points)-2, len(points)-1])

        # 水平线 (y = const)
        points.append([-half, i, 0])
        points.append([half, i, 0])
        lines.append([len(points)-2, len(points)-1])

    # 转为Open3D LineSet
    line_set = o3d.geometry.LineSet()
    line_set.points = o3d.utility.Vector3dVector(points)
    line_set.lines = o3d.utility.Vector2iVector(lines)
    line_set.colors = o3d.utility.Vector3dVector([[0.7, 0.7, 0.7] for _ in lines])  # 灰色

    return line_set


def visualize_trajectory_with_grid(poses, axis_size=0.1, grid_size=2.0, grid_interval=0.2):
    geometries = []

    # 坐标系们
    for pose in poses:
        x, y, z, qx, qy, qz, qw = pose
        rot = R.from_quat([qx, qy, qz, qw]).as_matrix()
        frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=axis_size)
        transform = np.eye(4)
        transform[:3, :3] = rot
        transform[:3, 3] = [x, y, z]
        frame.transform(transform)
        geometries.append(frame)

    # 轨迹连线
    points = [pose[:3] for pose in poses]
    lines = [[i, i+1] for i in range(len(points)-1)]
    line_set = o3d.geometry.LineSet(
        points=o3d.utility.Vector3dVector(points),
        lines=o3d.utility.Vector2iVector(lines),
    )
    line_set.colors = o3d.utility.Vector3dVector([[1, 0, 0] for _ in lines])  # 红色轨迹线
    geometries.append(line_set)

    # 添加grid
    grid = create_grid(grid_size, grid_interval)
    geometries.append(grid)

    o3d.visualization.draw_geometries(geometries)


# 示例
poses = [
    np.array([-0.025, 0, -0.035, 0, 0, 0, 1]),
    np.array([-0.02, 0, -0.035, 0, 0, 0, 1]),
    np.array([-0.015, 0, -0.035, 0, 0, 0, 1]),
    np.array([-0.01, 0., -0.03, 0, 0, 0, 1]),
    np.array([-0.005, 0., -0.025, 0, 0, 0, 1]),
    np.array([0.0, 0., -0.015, 0, 0, 0, 1]),
    np.array([0.005, 0., -0.005, 0, 0, 0, 1]),
    np.array([0.008, 0., 0.005, 0, 0, 0, 1]),
]

poses = [
    np.concatenate([np.array([0, 0, 0]), R.from_euler('xyz', [0.2, 1.5708, -0.8]).as_quat()]),
    np.concatenate([np.array([0.1, 0, 0]), R.from_euler('xyz', [1.0, 1.5708, 0.]).as_quat()]),
    np.concatenate([np.array([0.2, 0, 0]), R.from_euler('xyz', [0.4, 1.5708, -0.6]).as_quat()]),
    np.concatenate([np.array([0.3, 0, 0]), R.from_euler('xyz', [0.8, 1.5708, 0.]).as_quat()]),
    np.concatenate([np.array([0.4, 0, 0]), R.from_euler('xyz', [0.6, 1.5708, 0.]).as_quat()]),
]

visualize_trajectory_with_grid(poses, axis_size=0.05, grid_size=1., grid_interval=0.1)